/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
// ---------------------------------------------
// main.cpp ‒ test for OpenSpB_*cscmm
// ---------------------------------------------
#include "test_common.hpp"  
#include <iostream>
#include <vector>
#include <complex>
#include <random>
#include <cassert>
#include <cmath>

// -----------------------------------------------------------------------------
// 朴素 dense GEMM：C = beta*C + alpha*A*B
// A : m×k (列主序行距=1)
// B : k×n (列主序行距=1)
// C : m×n (列主序行距=ldc)
// -----------------------------------------------------------------------------
template <typename T>
void dense_gemm(int m, int n, int k,
                const T& alpha,
                const std::vector<T>& A,
                const std::vector<T>& B,
                const T& beta,
                std::vector<T>&       C, int ldc)
{
    // beta*C
    if (beta == T{}) {
        std::fill(C.begin(), C.end(), T{});
    } else if (beta != T{1}) {
        for (auto& x : C) x *= beta;
    }

    // alpha*A*B
    for (int colB = 0; colB < n; ++colB) {
        for (int rowA = 0; rowA < m; ++rowA) {
            T acc{};
            for (int p = 0; p < k; ++p)
                acc += A[rowA + p*m] * B[p + colB*k];
            C[rowA + colB*ldc] += alpha * acc;
        }
    }
}

// -----------------------------------------------------------------------------
// 随机生成器（仅用来给 C 初值，以确保 beta 路径被覆盖）
// -----------------------------------------------------------------------------
template <typename T>
T rand_val(std::mt19937& gen);
template<>
inline float rand_val<float>(std::mt19937& gen){
    static std::uniform_real_distribution<float> dist(-1.f,1.f);
    return dist(gen);
}
template<>
inline double rand_val<double>(std::mt19937& gen){
    static std::uniform_real_distribution<double> dist(-1.0,1.0);
    return dist(gen);
}
template<>
inline std::complex<float> rand_val<std::complex<float>>(std::mt19937& gen){
    return {rand_val<float>(gen), rand_val<float>(gen)};
}
template<>
inline std::complex<double> rand_val<std::complex<double>>(std::mt19937& gen){
    return {rand_val<double>(gen), rand_val<double>(gen)};
}

// -----------------------------------------------------------------------------
// 比较两块列主序矩阵是否近似相等
// -----------------------------------------------------------------------------
template <typename T>
bool near_equal(const std::vector<T>& ref,
                const std::vector<T>& got,
                double tol)
{
    double max_err = 0.0;
    for (size_t i = 0; i < ref.size(); ++i)
        max_err = std::max(max_err, static_cast<double>(std::abs(ref[i]-got[i])));
    return max_err <= tol;
}

// -----------------------------------------------------------------------------
// 构造一个 4×3 稀疏矩阵（6 个非零）及其 CSC 结构
//
// A = [ 1 0 2
//       0 3 0
//       4 0 5
//       0 6 0 ]
// -----------------------------------------------------------------------------
template <typename T>
void build_csc(std::vector<int>& colPtr,
               std::vector<int>& rowIdx,
               std::vector<T>&   valA)
{
    colPtr = {0,2,4,6};  // 0-based
    rowIdx = {0,2, 1,3, 0,2}; // 0-based
    valA   = {T{1}, T{4},
              T{3}, T{6},
              T{2}, T{5}};
}

// =============================================================================
// test for one-based indexing and beta=0
// =============================================================================
void test_one_based_beta_zero()
{
    using std::cout;
    std::mt19937 gen{456};

    // CSC格式矩阵A（1-based索引）
    // A = [1 4
    //      0 0
    //      2 5]
    std::vector<int> colPtr = {1, 3, 5}; // 1-based列指针
    std::vector<int> rowIdx = {1, 3, 1, 3}; // 1-based行索引
    OpenSpB_base_type  base = OPENSPBLAS_ONE_BASE;
    
    const int mA = 3, kA = 2; // A维度：3行2列
    const int n = 2;          // B和C的列数
    const int ldb = kA;       // B的行距
    const int ldc = mA;       // C的行距

    // -------------------------- 1. float / NO_TRANS --------------------------
    {
        using T = float;
        std::vector<T> valA = {1.0f, 2.0f, 4.0f, 5.0f};

        const int m = mA;
        const int k = kA;
        const T alpha = 1.0f;
        const T beta = 0.0f; // beta=0

        // B矩阵 (2×2): [1 3; 2 4]
        std::vector<T> B = {1.0f, 2.0f, 3.0f, 4.0f};
        // C初始化为非零值（但beta=0会清零）
        std::vector<T> C(m*n, 999.0f);

        // 参考结果: C = A*B
        // 第一列: [1 * 1 + 4 * 2 = 9, 0, 2 * 1 + 5 * 2 = 12]
        // 第二列: [1 * 3 + 4 * 4 = 19, 0, 2 * 3 + 5 * 4 = 26]
        std::vector<T> C_ref = {9.0f, 0.0f, 12.0f, 19.0f, 0.0f, 26.0f};

        OpenSpB_scscmm(OPENSPBLAS_NO_TRANS, base,
                          &m,&n,&k,
                          &alpha,
                          colPtr.data(), rowIdx.data(), valA.data(),
                          B.data(), ldb,
                          &beta,
                          C.data(), ldc);

        assert(near_equal(C_ref, C, 1e-6));
    }

    // -------------------------- 2. double / NO_TRANS -------------------------
    {
        using T = double;
        std::vector<T> valA = {1.0, 2.0, 4.0, 5.0};
        
        const int m = mA;
        const int k = kA;
        const T alpha = 1.0;
        const T beta = 0.0; // beta=0

        // B矩阵 (2×2)
        std::vector<T> B = {1.0, 2.0, 3.0, 4.0};
        std::vector<T> C(m*n, 999.0);
        
        // 参考结果
        std::vector<T> C_ref = {9.0, 0.0, 12.0, 19.0, 0.0, 26.0};

        OpenSpB_dcscmm(OPENSPBLAS_NO_TRANS, base,
                          &m,&n,&k,
                          &alpha,
                          colPtr.data(), rowIdx.data(), valA.data(),
                          B.data(), ldb,
                          &beta,
                          C.data(), ldc);

        assert(near_equal(C_ref, C, 1e-12));
    }

    // ------------------- 3. complex<float> / NO_TRANS -----------------------
    {
        using T = std::complex<float>;
        std::vector<T> valA = {{1,0},{2,0},{4,0},{5,0}};
        
        const int m = mA;
        const int k = kA;
        const T alpha{1.0f, 0.0f};
        const T beta{0.0f, 0.0f}; // beta=0+0i
        
        // B矩阵 (2×2): 添加虚部
        std::vector<T> B = {{1,1}, {2,0}, {3,-1}, {4,2}};
        std::vector<T> C(m*n, T{999,999});
        
        // 参考计算（复数计算）:
        // 第一列: [ (1,0)*(1,1) + (4,0)*(2,0) = (1,1) + (8,0) = (9,1)
        //          0
        //          (2,0)*(1,1) + (5,0)*(2,0) = (2,2) + (10,0) = (12,2) ]
        // 第二列: [ (1,0)*(3,-1) + (4,0)*(4,2) = (3,-1) + (16,8) = (19,7)
        //          0
        //          (2,0)*(3,-1) + (5,0)*(4,2) = (6,-2) + (20,10) = (26,8) ]
        std::vector<T> C_ref = { {9,1}, {0,0}, {12,2}, {19,7}, {0,0}, {26,8} };

        OpenSpB_ccscmm(OPENSPBLAS_NO_TRANS, base,
                          &m,&n,&k,
                          &alpha,
                          colPtr.data(), rowIdx.data(), valA.data(),
                          B.data(), ldb,
                          &beta,
                          C.data(), ldc);

        assert(near_equal(C_ref, C, 1e-5));
    }

    // ------------------- 4. complex<double> / NO_TRANS -----------------------
    {
        using T = std::complex<double>;
        std::vector<T> valA = {{1,0},{2,0},{4,0},{5,0}};
        
        const int m = mA;
        const int k = kA;
        const T alpha{1.0, 0.0};
        const T beta{0.0, 0.0}; // beta=0+0i
        
        // B矩阵 (2×2): 与complex<float>相同值
        std::vector<T> B = {{1,1}, {2,0}, {3,-1}, {4,2}};
        std::vector<T> C(m*n, T{999,999});
        
        // 参考结果（复数计算）
        std::vector<T> C_ref = { {9,1}, {0,0}, {12,2}, {19,7}, {0,0}, {26,8} };

        OpenSpB_zcscmm(OPENSPBLAS_NO_TRANS, base,
                          &m,&n,&k,
                          &alpha,
                          colPtr.data(), rowIdx.data(), valA.data(),
                          B.data(), ldb,
                          &beta,
                          C.data(), ldc);

        assert(near_equal(C_ref, C, 1e-12));
    }

    std::cout << "One-based index and beta=0 tests passed!\n";
}

// =============================================================================
// main
// =============================================================================
TEST(SparseBLASlevel2, spmm) {
    using std::cout;
    std::mt19937 gen{123};

    // common geometry ---------------------------------------------------------
    const int mA = 4, kA = 3;
    OpenSpB_base_type  base = OPENSPBLAS_ZERO_BASE;
    std::vector<int> colPtr, rowIdx;

    // -------------------------- 1. float / NO_TRANS --------------------------
    {
        using T = float;
        std::vector<T> valA;
        build_csc(colPtr,rowIdx,valA);

        const int m = mA;
        const int k = kA;
        const int n = 2;
        const T alpha = 2.0f, beta = 0.5f;
        const int ldb = k;
        const int ldc = m;

        // B
        std::vector<T> B = {7,8,9, 1,2,3};   // k×n 列主序
        // C (随机初始化，确保 beta 路径生效)
        std::vector<T> C(m*n);
        for (auto& x:C) x = rand_val<T>(gen);

        // reference
        auto C_ref = C;
        std::vector<T> A_dense(m*k);
        // 生成 dense A (列主序)
        for (int col=0; col<k; ++col){
            for (int p=colPtr[col]; p<colPtr[col+1]; ++p){
                int row = rowIdx[p];
                A_dense[row + col*m] = valA[p];
            }
        }
        dense_gemm(m,n,k,alpha,A_dense,B,beta,C_ref,ldc);

        // spMM
        OpenSpB_scscmm(OPENSPBLAS_NO_TRANS, base,
                          &m,&n,&k,
                          &alpha,
                          colPtr.data(),rowIdx.data(),valA.data(),
                          B.data(),ldb,
                          &beta,
                          C.data(),ldc);

        assert(near_equal(C_ref,C,1e-6));
    }

    // -------------------------- 2. double / TRANS ---------------------------
    {
        using T = double;
        std::vector<T> valA;
        build_csc(colPtr,rowIdx,valA);

        const int m = kA;   // m = kA after transpose
        const int k = mA;
        const int n = 2;
        const T alpha = 2.0, beta = 0.5;
        const int ldb = k;
        const int ldc = m;

        // B : k×n
        std::vector<T> B = {4,5,6,7, 0,1,2,3};  // 4×2
        std::vector<T> C(m*n);
        for (auto& x:C) x = rand_val<T>(gen);

        // reference
        auto C_ref = C;
        // 生成 dense Aᵀ
        std::vector<T> AT_dense(m*k);        // 3×4 转置 → 3(cols)? 正确尺寸:3x4
        for (int col=0; col<kA; ++col){
            for (int p=colPtr[col]; p<colPtr[col+1]; ++p){
                int row = rowIdx[p];
                AT_dense[col + row*m] = valA[p];
            }
        }
        dense_gemm(m,n,k,alpha,AT_dense,B,beta,C_ref,ldc);

        // spMM
        OpenSpB_dcscmm(OPENSPBLAS_TRANS, base,
                          &m,&n,&k,
                          &alpha,
                          colPtr.data(),rowIdx.data(),valA.data(),
                          B.data(),ldb,
                          &beta,
                          C.data(),ldc);

        assert(near_equal(C_ref,C,1e-12));
    }

    // ------------------- 3. complex<float> / NO_TRANS -----------------------
    {
        using T = std::complex<float>;
        std::vector<T> valA;
        build_csc(colPtr,rowIdx,valA);

        const int m=mA,k=kA,n=2;
        const T alpha = T{1.5f,0.7f};
        const T beta  = T{-0.3f,0.2f};
        const int ldb=k, ldc=m;

        std::vector<T> B = { {1,1},{2,0}, {3,-1}, {4,2}, {5,0},{6,-2} };
        std::vector<T> C(m*n);
        for (auto& x:C) x = rand_val<T>(gen);

        auto C_ref=C;
        std::vector<T> A_dense(m*k);
        for(int col=0; col<k; ++col)
            for(int p=colPtr[col]; p<colPtr[col+1]; ++p)
                A_dense[rowIdx[p]+col*m]=valA[p];

        dense_gemm(m,n,k,alpha,A_dense,B,beta,C_ref,ldc);

        OpenSpB_ccscmm(OPENSPBLAS_NO_TRANS, base,
                          &m,&n,&k,
                          &alpha,
                          colPtr.data(),rowIdx.data(),valA.data(),
                          B.data(),ldb,
                          &beta,
                          C.data(),ldc);

        assert(near_equal(C_ref,C,1e-5));
    }

    // ------------- 4. complex<double> / CONJ_TRANS --------------------------
    {
        using T = std::complex<double>;
        std::vector<T> valA;
        build_csc(colPtr,rowIdx,valA);

        const int m=kA,k=mA,n=2;
        const T alpha = {0.8, -1.1};
        const T beta  = {0.0, 0.2};
        const int ldb=k, ldc=m;

        std::vector<T> B = { {0,2},{1,1},{2,0},{3,-1}, {4,0},{5,2},{6,1},{7,-3} };
        std::vector<T> C(m*n);
        for (auto& x:C) x = rand_val<T>(gen);

        auto C_ref=C;
        // 生成 conj(A)ᵀ = Aᴴ
        std::vector<T> AH_dense(m*k);   // 3×4
        for(int col=0; col<kA; ++col)
            for(int p=colPtr[col]; p<colPtr[col+1]; ++p)
                AH_dense[col + rowIdx[p]*m]=std::conj(valA[p]);

        dense_gemm(m,n,k,alpha,AH_dense,B,beta,C_ref,ldc);

        OpenSpB_zcscmm(OPENSPBLAS_CONJ_TRANS, base,
                          &m,&n,&k,
                          &alpha,
                          colPtr.data(),rowIdx.data(),valA.data(),
                          B.data(),ldb,
                          &beta,
                          C.data(),ldc);

        assert(near_equal(C_ref,C,1e-12));
    }

    // 运行新增的1-based和beta=0测试
    test_one_based_beta_zero();

    cout << "All tests passed!\n";
}